{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e1bc453d-c8d3-4503-b3da-52120ad92c74",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Reflection Pattern\n",
    "\n",
    "source : https://github.com/neural-maze/agentic_patterns\n",
    "\n",
    "The first pattern we are going to implement is the **reflection pattern**. \n",
    "\n",
    "\n",
    "This pattern allows the LLM to reflect and critique its outputs, following the next steps:\n",
    "\n",
    "1. The LLM **generates** a candidate output. If you look at the diagram above, it happens inside the **\"Generate\"** box.\n",
    "2. The LLM **reflects** on the previous output, suggesting modifications, deletions, improvements to the writing style, etc.\n",
    "3. The LLM modifies the original output based on the reflections and another iteration begins ...\n",
    "\n",
    "**Now, we are going to build, from scratch, each step, so that you can truly understand how this pattern works.**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7898c34d-de9a-4970-b7f4-3d86b69d45a7",
   "metadata": {},
   "source": [
    "## Generation Step"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "031f6b07-4f99-46f6-a53c-ff242585cbad",
   "metadata": {},
   "source": [
    "The first thing we need to consider is:\n",
    "\n",
    "> What do we want to generate? A poem? An essay? Python code?\n",
    "\n",
    "For this example, I've decided to test the Python coding skills of Llama3 70B (that's the LLM we are going to use for all the tutorials). In particular, we are going to ask our LLM to code a famous sorting algorithm: **Merge Sort**. \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73f4d7b7-40bf-43b9-a626-2a11d5529ac8",
   "metadata": {},
   "source": [
    "### Groq Client and relevant imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96731d2f-a079-4e41-9756-220f02d4ebd8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from pprint import pprint\n",
    "from groq import Groq\n",
    "from dotenv import load_dotenv\n",
    "from IPython.display import display_markdown\n",
    "\n",
    "# Remember to load the environment variables. You should have the Groq API Key in there :)\n",
    "load_dotenv()\n",
    "\n",
    "client = Groq()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e644a635-e035-44e2-8c25-cee0f2b56556",
   "metadata": {},
   "source": [
    "We will start the **\"generation\"** chat history with the system prompt, as we said before. In this case, let the LLM act like a Python \n",
    "programmer eager to receive feedback / critique by the user."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12467256-c741-495a-9923-439c1fcf270d",
   "metadata": {},
   "outputs": [],
   "source": [
    "generation_chat_history = [\n",
    "    {\n",
    "        \"role\": \"system\",\n",
    "        \"content\": \"You are a Python programmer tasked with generating high quality Python code.\"\n",
    "        \"Your task is to Generate the best content possible for the user's request. If the user provides critique,\" \n",
    "        \"respond with a revised version of your previous attempt.\"\n",
    "    }\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43149b4f-54db-455f-9d39-6ad2f5c52b94",
   "metadata": {},
   "source": [
    "Now, as the user, we are going to ask the LLM to generate an implementation of the **Merge Sort** algorithm. Just add a new message with the **user** role to the chat history."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0742e7bd-4857-4ed1-a96b-37098d448bdd",
   "metadata": {},
   "outputs": [],
   "source": [
    "generation_chat_history.append(\n",
    "    {\n",
    "        \"role\": \"user\",\n",
    "        \"content\": \"Generate a Python implementation of the Merge Sort algorithm\"\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4df1bffe-375f-4a9a-8433-e217eb94aea2",
   "metadata": {},
   "source": [
    "Let's generate the first version of the essay."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff984277-733c-4495-b7fd-0669393380b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "mergesort_code = client.chat.completions.create(\n",
    "    messages=generation_chat_history,\n",
    "    model=\"llama3-70b-8192\"\n",
    ").choices[0].message.content\n",
    "\n",
    "generation_chat_history.append(\n",
    "    {\n",
    "        \"role\": \"assistant\",\n",
    "        \"content\": mergesort_code\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c03f208b-2234-4fd1-a02b-f4fff06c01a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "display_markdown(mergesort_code, raw=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a04ebe5-0573-4520-a529-aff22d486b7d",
   "metadata": {},
   "source": [
    "## Reflection Step"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67aa69e4-632f-4a0c-a6f0-c5a7ced4849d",
   "metadata": {
    "tags": []
   },
   "source": [
    "Now, let's allow the LLM to reflect on its outputs by defining another system prompt. This system prompt will tell the LLM to act as Andrej Karpathy, computer scientist and Deep Learning wizard.\n",
    "\n",
    ">To be honest, I don't think the fact of acting like Andrej Karpathy will influence the LLM outputs, but it was fun :)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d93c928-d585-48af-a74c-a5b8d84593c6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "reflection_chat_history = [\n",
    "    {\n",
    "    \"role\": \"system\",\n",
    "    \"content\": \"You are Andrej Karpathy, an experienced computer scientist. You are tasked with generating critique and recommendations for the user's code\",\n",
    "    }\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c498175f-b3f9-40af-92a3-d5b36d77d1cf",
   "metadata": {},
   "source": [
    "The user message, in this case,  is the essay generated in the previous step. We simply add the `mergesort_code` to the `reflection_chat_history`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26af1a73-4d91-40e8-a9bc-c34d32b2ab82",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "reflection_chat_history.append(\n",
    "    {\n",
    "        \"role\": \"user\",\n",
    "        \"content\": mergesort_code\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bfa994c8-3612-47b0-9571-e21d0d73d896",
   "metadata": {},
   "source": [
    "Now, let's generate a critique to the Python code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40fee42f-d47a-41b1-a40d-7208ba76ce98",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "critique = client.chat.completions.create(\n",
    "    messages=reflection_chat_history,\n",
    "    model=\"llama3-70b-8192\"\n",
    ").choices[0].message.content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0fef3203-c7f1-407f-8b9b-4e8ae140a4cb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "display_markdown(critique, raw=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5df433b0-d662-4378-895e-6b09dd3201bc",
   "metadata": {},
   "source": [
    "Finally, we just need to add this *critique* to the `generation_chat_history`, in this case, as the `user` role."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27a85bb3-cf6a-4576-8caf-cd41e602a1f1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "generation_chat_history.append(\n",
    "    {\n",
    "        \"role\": \"user\",\n",
    "        \"content\": critique\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3c1aefa-8454-41ab-af40-2675f340a577",
   "metadata": {},
   "source": [
    "## Generation Step (II)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91d845cf-51c3-4cfd-b6a7-1b970413f6db",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "essay = client.chat.completions.create(\n",
    "    messages=generation_chat_history,\n",
    "    model=\"llama3-70b-8192\"\n",
    ").choices[0].message.content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef14eaa8-f501-4efc-997f-8564ec8dccd8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "display_markdown(essay, raw=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75883af2-f31d-4c24-b1ff-315a0711f9fa",
   "metadata": {},
   "source": [
    "## And the iteration starts again ..."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5b824d1-c17e-448c-bdd7-df543aa5a9fd",
   "metadata": {},
   "source": [
    "After **Generation Step (II)** the corrected Python code will be received, once again, by Karpathy. Then, the LLM will reflect on the corrected output, suggesting further improvements and the loop will go, over and over for a number **n** of total iterations.\n",
    "\n",
    "> There's another possibility. Suppose the Reflection step can't find any further improvement. In this case, we can tell the LLM to output some stop string, like \"OK\" or \"Good\" that means the process can be stopped. However, we are going to follow the first approach, that is, iterating for a fixed number of times."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7cf2cf5b-d083-435c-914a-3ff484d53473",
   "metadata": {},
   "source": [
    "## Implementing a class "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15f9a9e6-29f3-4adf-863e-c49fbb9a6b44",
   "metadata": {},
   "source": [
    "Now that you understand the underlying loop of the Reflection Agent, let's implement this agent as a class."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34ea5aaf",
   "metadata": {},
   "source": [
    "#### utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "987f3015",
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "from dataclasses import dataclass\n",
    "import time\n",
    "\n",
    "from colorama import Fore\n",
    "from colorama import Style\n",
    "\n",
    "def completions_create(client, messages: list, model: str) -> str:\n",
    "    \"\"\"\n",
    "    Sends a request to the client's `completions.create` method to interact with the language model.\n",
    "\n",
    "    Args:\n",
    "        client (Groq): The Groq client object\n",
    "        messages (list[dict]): A list of message objects containing chat history for the model.\n",
    "        model (str): The model to use for generating tool calls and responses.\n",
    "\n",
    "    Returns:\n",
    "        str: The content of the model's response.\n",
    "    \"\"\"\n",
    "    response = client.chat.completions.create(messages=messages, model=model)\n",
    "    return str(response.choices[0].message.content)\n",
    "\n",
    "\n",
    "def build_prompt_structure(prompt: str, role: str, tag: str = \"\") -> dict:\n",
    "    \"\"\"\n",
    "    Builds a structured prompt that includes the role and content.\n",
    "\n",
    "    Args:\n",
    "        prompt (str): The actual content of the prompt.\n",
    "        role (str): The role of the speaker (e.g., user, assistant).\n",
    "\n",
    "    Returns:\n",
    "        dict: A dictionary representing the structured prompt.\n",
    "    \"\"\"\n",
    "    if tag:\n",
    "        prompt = f\"<{tag}>{prompt}</{tag}>\"\n",
    "    return {\"role\": role, \"content\": prompt}\n",
    "\n",
    "\n",
    "def update_chat_history(history: list, msg: str, role: str):\n",
    "    \"\"\"\n",
    "    Updates the chat history by appending the latest response.\n",
    "\n",
    "    Args:\n",
    "        history (list): The list representing the current chat history.\n",
    "        msg (str): The message to append.\n",
    "        role (str): The role type (e.g. 'user', 'assistant', 'system')\n",
    "    \"\"\"\n",
    "    history.append(build_prompt_structure(prompt=msg, role=role))\n",
    "\n",
    "\n",
    "class ChatHistory(list):\n",
    "    def __init__(self, messages: list | None = None, total_length: int = -1):\n",
    "        \"\"\"Initialise the queue with a fixed total length.\n",
    "\n",
    "        Args:\n",
    "            messages (list | None): A list of initial messages\n",
    "            total_length (int): The maximum number of messages the chat history can hold.\n",
    "        \"\"\"\n",
    "        if messages is None:\n",
    "            messages = []\n",
    "\n",
    "        super().__init__(messages)\n",
    "        self.total_length = total_length\n",
    "\n",
    "    def append(self, msg: str):\n",
    "        \"\"\"Add a message to the queue.\n",
    "\n",
    "        Args:\n",
    "            msg (str): The message to be added to the queue\n",
    "        \"\"\"\n",
    "        if len(self) == self.total_length:\n",
    "            self.pop(0)\n",
    "        super().append(msg)\n",
    "\n",
    "\n",
    "class FixedFirstChatHistory(ChatHistory):\n",
    "    def __init__(self, messages: list | None = None, total_length: int = -1):\n",
    "        \"\"\"Initialise the queue with a fixed total length.\n",
    "\n",
    "        Args:\n",
    "            messages (list | None): A list of initial messages\n",
    "            total_length (int): The maximum number of messages the chat history can hold.\n",
    "        \"\"\"\n",
    "        super().__init__(messages, total_length)\n",
    "\n",
    "    def append(self, msg: str):\n",
    "        \"\"\"Add a message to the queue. The first messaage will always stay fixed.\n",
    "\n",
    "        Args:\n",
    "            msg (str): The message to be added to the queue\n",
    "        \"\"\"\n",
    "        if len(self) == self.total_length:\n",
    "            self.pop(1)\n",
    "        super().append(msg)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class TagContentResult:\n",
    "    \"\"\"\n",
    "    A data class to represent the result of extracting tag content.\n",
    "\n",
    "    Attributes:\n",
    "        content (List[str]): A list of strings containing the content found between the specified tags.\n",
    "        found (bool): A flag indicating whether any content was found for the given tag.\n",
    "    \"\"\"\n",
    "\n",
    "    content: list[str]\n",
    "    found: bool\n",
    "\n",
    "\n",
    "def extract_tag_content(text: str, tag: str) -> TagContentResult:\n",
    "    \"\"\"\n",
    "    Extracts all content enclosed by specified tags (e.g., <thought>, <response>, etc.).\n",
    "\n",
    "    Parameters:\n",
    "        text (str): The input string containing multiple potential tags.\n",
    "        tag (str): The name of the tag to search for (e.g., 'thought', 'response').\n",
    "\n",
    "    Returns:\n",
    "        dict: A dictionary with the following keys:\n",
    "            - 'content' (list): A list of strings containing the content found between the specified tags.\n",
    "            - 'found' (bool): A flag indicating whether any content was found for the given tag.\n",
    "    \"\"\"\n",
    "    # Build the regex pattern dynamically to find multiple occurrences of the tag\n",
    "    tag_pattern = rf\"<{tag}>(.*?)</{tag}>\"\n",
    "\n",
    "    # Use findall to capture all content between the specified tag\n",
    "    matched_contents = re.findall(tag_pattern, text, re.DOTALL)\n",
    "\n",
    "    # Return the dataclass instance with the result\n",
    "    return TagContentResult(\n",
    "        content=[content.strip() for content in matched_contents],\n",
    "        found=bool(matched_contents),\n",
    "    )\n",
    "\n",
    "\n",
    "def fancy_print(message: str) -> None:\n",
    "    \"\"\"\n",
    "    Displays a fancy print message.\n",
    "\n",
    "    Args:\n",
    "        message (str): The message to display.\n",
    "    \"\"\"\n",
    "    print(Style.BRIGHT + Fore.CYAN + f\"\\n{'=' * 50}\")\n",
    "    print(Fore.MAGENTA + f\"{message}\")\n",
    "    print(Style.BRIGHT + Fore.CYAN + f\"{'=' * 50}\\n\")\n",
    "    time.sleep(0.5)\n",
    "\n",
    "\n",
    "def fancy_step_tracker(step: int, total_steps: int) -> None:\n",
    "    \"\"\"\n",
    "    Displays a fancy step tracker for each iteration of the generation-reflection loop.\n",
    "\n",
    "    Args:\n",
    "        step (int): The current step in the loop.\n",
    "        total_steps (int): The total number of steps in the loop.\n",
    "    \"\"\"\n",
    "    fancy_print(f\"STEP {step + 1}/{total_steps}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f904241-29a1-4519-b6ab-15be0a7cfc53",
   "metadata": {},
   "outputs": [],
   "source": [
    "from colorama import Fore\n",
    "from dotenv import load_dotenv\n",
    "from groq import Groq\n",
    "load_dotenv()\n",
    "\n",
    "\n",
    "BASE_GENERATION_SYSTEM_PROMPT = \"\"\"\n",
    "Your task is to Generate the best content possible for the user's request.\n",
    "If the user provides critique, respond with a revised version of your previous attempt.\n",
    "You must always output the revised content.\n",
    "\"\"\"\n",
    "\n",
    "BASE_REFLECTION_SYSTEM_PROMPT = \"\"\"\n",
    "You are tasked with generating critique and recommendations to the user's generated content.\n",
    "If the user content has something wrong or something to be improved, output a list of recommendations\n",
    "and critiques. If the user content is ok and there's nothing to change, output this: <OK>\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "class ReflectionAgent:\n",
    "    \"\"\"\n",
    "    A class that implements a Reflection Agent, which generates responses and reflects\n",
    "    on them using the LLM to iteratively improve the interaction. The agent first generates\n",
    "    responses based on provided prompts and then critiques them in a reflection step.\n",
    "\n",
    "    Attributes:\n",
    "        model (str): The model name used for generating and reflecting on responses.\n",
    "        client (Groq): An instance of the Groq client to interact with the language model.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, model: str = \"llama-3.1-70b-versatile\"):\n",
    "        self.client = Groq()\n",
    "        self.model = model\n",
    "\n",
    "    def _request_completion(\n",
    "        self,\n",
    "        history: list,\n",
    "        verbose: int = 0,\n",
    "        log_title: str = \"COMPLETION\",\n",
    "        log_color: str = \"\",\n",
    "    ):\n",
    "        \"\"\"\n",
    "        A private method to request a completion from the Groq model.\n",
    "\n",
    "        Args:\n",
    "            history (list): A list of messages forming the conversation or reflection history.\n",
    "            verbose (int, optional): The verbosity level. Defaults to 0 (no output).\n",
    "\n",
    "        Returns:\n",
    "            str: The model-generated response.\n",
    "        \"\"\"\n",
    "        output = completions_create(self.client, history, self.model)\n",
    "\n",
    "        if verbose > 0:\n",
    "            print(log_color, f\"\\n\\n{log_title}\\n\\n\", output)\n",
    "\n",
    "        return output\n",
    "\n",
    "    def generate(self, generation_history: list, verbose: int = 0) -> str:\n",
    "        \"\"\"\n",
    "        Generates a response based on the provided generation history using the model.\n",
    "\n",
    "        Args:\n",
    "            generation_history (list): A list of messages forming the conversation or generation history.\n",
    "            verbose (int, optional): The verbosity level, controlling printed output. Defaults to 0.\n",
    "\n",
    "        Returns:\n",
    "            str: The generated response.\n",
    "        \"\"\"\n",
    "        return self._request_completion(\n",
    "            generation_history, verbose, log_title=\"GENERATION\", log_color=Fore.BLUE\n",
    "        )\n",
    "\n",
    "    def reflect(self, reflection_history: list, verbose: int = 0) -> str:\n",
    "        \"\"\"\n",
    "        Reflects on the generation history by generating a critique or feedback.\n",
    "\n",
    "        Args:\n",
    "            reflection_history (list): A list of messages forming the reflection history, typically based on\n",
    "                                       the previous generation or interaction.\n",
    "            verbose (int, optional): The verbosity level, controlling printed output. Defaults to 0.\n",
    "\n",
    "        Returns:\n",
    "            str: The critique or reflection response from the model.\n",
    "        \"\"\"\n",
    "        return self._request_completion(\n",
    "            reflection_history, verbose, log_title=\"REFLECTION\", log_color=Fore.GREEN\n",
    "        )\n",
    "\n",
    "    def run(\n",
    "        self,\n",
    "        user_msg: str,\n",
    "        generation_system_prompt: str = \"\",\n",
    "        reflection_system_prompt: str = \"\",\n",
    "        n_steps: int = 10,\n",
    "        verbose: int = 0,\n",
    "    ) -> str:\n",
    "        \"\"\"\n",
    "        Runs the ReflectionAgent over multiple steps, alternating between generating a response\n",
    "        and reflecting on it for the specified number of steps.\n",
    "\n",
    "        Args:\n",
    "            user_msg (str): The user message or query that initiates the interaction.\n",
    "            generation_system_prompt (str, optional): The system prompt for guiding the generation process.\n",
    "            reflection_system_prompt (str, optional): The system prompt for guiding the reflection process.\n",
    "            n_steps (int, optional): The number of generate-reflect cycles to perform. Defaults to 3.\n",
    "            verbose (int, optional): The verbosity level controlling printed output. Defaults to 0.\n",
    "\n",
    "        Returns:\n",
    "            str: The final generated response after all cycles are completed.\n",
    "        \"\"\"\n",
    "        generation_system_prompt += BASE_GENERATION_SYSTEM_PROMPT\n",
    "        reflection_system_prompt += BASE_REFLECTION_SYSTEM_PROMPT\n",
    "\n",
    "        # Given the iterative nature of the Reflection Pattern, we might exhaust the LLM context (or\n",
    "        # make it really slow). That's the reason I'm limitting the chat history to three messages.\n",
    "        # The `FixedFirstChatHistory` is a very simple class, that creates a Queue that always keeps\n",
    "        # fixeed the first message. I thought this would be useful for maintaining the system prompt\n",
    "        # in the chat history.\n",
    "        generation_history = FixedFirstChatHistory(\n",
    "            [\n",
    "                build_prompt_structure(prompt=generation_system_prompt, role=\"system\"),\n",
    "                build_prompt_structure(prompt=user_msg, role=\"user\"),\n",
    "            ],\n",
    "            total_length=3,\n",
    "        )\n",
    "\n",
    "        reflection_history = FixedFirstChatHistory(\n",
    "            [build_prompt_structure(prompt=reflection_system_prompt, role=\"system\")],\n",
    "            total_length=3,\n",
    "        )\n",
    "\n",
    "        for step in range(n_steps):\n",
    "            if verbose > 0:\n",
    "                fancy_step_tracker(step, n_steps)\n",
    "\n",
    "            # Generate the response\n",
    "            generation = self.generate(generation_history, verbose=verbose)\n",
    "            update_chat_history(generation_history, generation, \"assistant\")\n",
    "            update_chat_history(reflection_history, generation, \"user\")\n",
    "\n",
    "            # Reflect and critique the generation\n",
    "            critique = self.reflect(reflection_history, verbose=verbose)\n",
    "\n",
    "            if \"<OK>\" in critique:\n",
    "                # If no additional suggestions are made, stop the loop\n",
    "                print(\n",
    "                    Fore.RED,\n",
    "                    \"\\n\\nStop Sequence found. Stopping the reflection loop ... \\n\\n\",\n",
    "                )\n",
    "                break\n",
    "\n",
    "            update_chat_history(generation_history, critique, \"user\")\n",
    "            update_chat_history(reflection_history, critique, \"assistant\")\n",
    "\n",
    "        return generation\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd1a8071-c763-4dbf-8db7-60f9116f62e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "agent = ReflectionAgent()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87c8cf16-0dfa-49b6-bc30-8f14bbe7860a",
   "metadata": {},
   "outputs": [],
   "source": [
    "generation_system_prompt = \"You are a Python programmer tasked with generating high quality Python code\"\n",
    "\n",
    "reflection_system_prompt = \"You are Andrej Karpathy, an experienced computer scientist\"\n",
    "\n",
    "user_msg = \"Generate a Python implementation of the Merge Sort algorithm\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a9a3e5b-9b45-4a27-b391-f78b57ff94f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "final_response = agent.run(\n",
    "    user_msg=user_msg,\n",
    "    generation_system_prompt=generation_system_prompt,\n",
    "    reflection_system_prompt=reflection_system_prompt,\n",
    "    n_steps=10,\n",
    "    verbose=1,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b69d182-d12e-40bb-8dfb-cbc8903218a1",
   "metadata": {},
   "source": [
    "## Final result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e4663cd-61dd-4a38-866a-f032045a444a",
   "metadata": {},
   "outputs": [],
   "source": [
    "display_markdown(final_response, raw=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
